#define JZQUTILS
#define GEOMIPMAP

#include <glew.h>

#include <cmath>
#include <algorithm>
#include <limits>
#include <vector>

#include <cstdio>

using namespace std;

#ifdef JZQUTILS
  #include "Vec.h"
  #include "Mat.h"
  #include "Array2.h"
#endif

#ifdef GEOMIPMAP
  #include "geomipmap.h"
#endif

namespace
{
#ifndef JZQUTILS
  ///  
#endif

#ifndef GEOMIPMAP
  ///
#endif

  class GLHeightField
  {
  public:
    GLHeightField()
    {
      
    }
    
    ~GLHeightField()
    {
      
    }
    
    void init(int width,int depth,int chunkSize)
    {
      assert( ((width-1)%(chunkSize-1)) == 0 );
      assert( ((depth-1)%(chunkSize-1)) == 0 );
  
      this->width = width;
      this->depth = depth;
      this->chunkSize = chunkSize;
      
      chunksPerWidth = (width-1) / (chunkSize-1);
      chunksPerDepth = (depth-1) / (chunkSize-1);
      
      maxDim = max(width,depth);
        
      numLods = geomipmap_num_lods(chunkSize);
            
      chunkIbo = Array2<GLuint>(numLods,16);
    
      for(int lod=0;lod<numLods;lod++)
      {
        for(int stitchMask=0;stitchMask<16;stitchMask++)
        {
          int numIds = geomipmap_num_ids(chunkSize,lod,stitchMask);
      
          int* ids = new int[numIds];
          
          geomipmap_gen_ids(chunkSize,lod,stitchMask,chunkSize,ids);
      
          glGenBuffers(1,&chunkIbo(lod,stitchMask));
          glBindBuffer(GL_ELEMENT_ARRAY_BUFFER,chunkIbo(lod,stitchMask));
          glBufferData(GL_ELEMENT_ARRAY_BUFFER,numIds*sizeof(int),ids,GL_STATIC_DRAW);
      
          delete[] ids;
        }
      }
   
   
      glGenBuffers(1,&vbo);
      glBindBuffer(GL_ARRAY_BUFFER,vbo);
      
      int heightFieldBufferSize = (chunksPerWidth*chunkSize)*(chunksPerDepth*chunkSize)*2*sizeof(float);
      int borderMaskBufferSize = chunkSize*chunkSize*4*sizeof(float);
      int vertexLevelBufferSize = chunkSize*chunkSize*sizeof(float);
      int vertexCoordBufferSize = chunkSize*chunkSize*2*sizeof(float);
      
      int heightFieldBufferOffset = 0;
      borderMaskBufferOffset = heightFieldBufferOffset+heightFieldBufferSize;
      vertexLevelBufferOffset = borderMaskBufferOffset+borderMaskBufferSize;
      vertexCoordBufferOffset = vertexLevelBufferOffset+vertexLevelBufferSize;
      
      glBufferData(GL_ARRAY_BUFFER,heightFieldBufferSize+borderMaskBufferSize+vertexLevelBufferSize+vertexCoordBufferSize, NULL, GL_DYNAMIC_DRAW);
    
      float* borderMask = new float[chunkSize*chunkSize*4];
    
      for(int j=0;j<chunkSize;j++)
      for(int i=0;i<chunkSize;i++)
      {
        borderMask[(i+j*chunkSize)*4+0] = (i==0           && (j>0 && j<chunkSize-1)) ? 1.0f : 0.0f;
        borderMask[(i+j*chunkSize)*4+1] = (i==chunkSize-1 && (j>0 && j<chunkSize-1)) ? 1.0f : 0.0f;
        borderMask[(i+j*chunkSize)*4+2] = (j==0           && (i>0 && i<chunkSize-1)) ? 1.0f : 0.0f;
        borderMask[(i+j*chunkSize)*4+3] = (j==chunkSize-1 && (i>0 && i<chunkSize-1)) ? 1.0f : 0.0f;
      }
    
      glBufferSubData(GL_ARRAY_BUFFER,borderMaskBufferOffset,borderMaskBufferSize,borderMask);
    
      delete[] borderMask;
      
      float* vertexLevel = new float[chunkSize*chunkSize];
  
    
      for(int j=0;j<chunkSize;j++)
      for(int i=0;i<chunkSize;i++)
      {
        int level = 0;
    
        for(int lod=1;lod<numLods;lod++)
        {
          int pitch = 1 << lod;
          if ((i%pitch)==0 && (j%pitch)==0) level = lod;
        }
  
        vertexLevel[(i+j*chunkSize)] = level;
      }
    
      glBufferSubData(GL_ARRAY_BUFFER,vertexLevelBufferOffset,vertexLevelBufferSize,vertexLevel);
    
      delete[] vertexLevel;
      
      float* vertexCoord = new float[chunkSize*chunkSize*2];
    
      for(int j=0;j<chunkSize;j++)
      for(int i=0;i<chunkSize;i++)
      {
        vertexCoord[(i+j*chunkSize)*2+0] = float(i)/float(maxDim-1);
        vertexCoord[(i+j*chunkSize)*2+1] = float(j)/float(maxDim-1);
      }
      
      glBufferSubData(GL_ARRAY_BUFFER,vertexCoordBufferOffset,vertexCoordBufferSize,vertexCoord);
    
      delete[] vertexCoord;
      
      chunkCenter = Array2<Vec3f>(chunksPerWidth,chunksPerDepth);
  
      chunkMaxLodDist = new Array2<float>[numLods];
  
      for(int lod=0;lod<numLods;lod++)
      {
        chunkMaxLodDist[lod] = Array2<float>(chunksPerWidth,chunksPerDepth);
      }
      
      chunkLod = Array2<float>(chunksPerWidth,chunksPerDepth);
      chunkMinMaxHeight = Array2<Vec2f>(chunksPerWidth,chunksPerDepth);
  
      heightFieldBuffer = new float[(chunksPerWidth*chunkSize)*(chunksPerDepth*chunkSize)*2];
  
      offsetTable = new int[chunkSize*chunkSize*2];
  
      for(int vj=0;vj<chunkSize;vj++)
      {
        for(int vi=0;vi<chunkSize;vi++)
        {
          int level = 0;
                     
          for(int lod=1;lod<numLods;lod++)
          {
            int pitch = 1 << lod;
            if ((vi%pitch)==0 && (vj%pitch)==0) level = lod;                            
          }
          
          int pitch = 1 << level;
          
          int pi = (vi / pitch)%4;
          int pj = (vj / pitch)%4;
          
          if (level<numLods-1)
          {
            if (pj%2==0)
            {
              offsetTable[(vi+vj*chunkSize)*2+0] = -pitch;
              offsetTable[(vi+vj*chunkSize)*2+1] = +pitch;
            }
            else if (pi%2==0)
            {
              offsetTable[(vi+vj*chunkSize)*2+0] = -pitch*width;
              offsetTable[(vi+vj*chunkSize)*2+1] = +pitch*width;
            }
            else
            {        
              if (level==numLods-2)
              {
                offsetTable[(vi+vj*chunkSize)*2+0] = -pitch-pitch*width;
                offsetTable[(vi+vj*chunkSize)*2+1] = +pitch+pitch*width;
              }
              else
              {
                if      (pi==1 && pj==1) { offsetTable[(vi+vj*chunkSize)*2+0] = -pitch-pitch*width; offsetTable[(vi+vj*chunkSize)*2+1] = +pitch+pitch*width; }
                else if (pi==3 && pj==1) { offsetTable[(vi+vj*chunkSize)*2+0] = +pitch-pitch*width; offsetTable[(vi+vj*chunkSize)*2+1] = -pitch+pitch*width; }
                else if (pi==1 && pj==3) { offsetTable[(vi+vj*chunkSize)*2+0] = -pitch+pitch*width; offsetTable[(vi+vj*chunkSize)*2+1] = +pitch-pitch*width; }
                else if (pi==3 && pj==3) { offsetTable[(vi+vj*chunkSize)*2+0] = +pitch+pitch*width; offsetTable[(vi+vj*chunkSize)*2+1] = -pitch-pitch*width; }
              }
            }
          }
          else
          {
            offsetTable[(vi+vj*chunkSize)*2+0] = 0;
            offsetTable[(vi+vj*chunkSize)*2+1] = 0;
          }
        }      
      }  
    }
    
    void update(const float* heightData)
    {
      int heightFieldBufferOffset = 0;
      
      for(int cj=0;cj<chunksPerDepth;cj++)
      {
        for(int ci=0;ci<chunksPerWidth;ci++)
        {
          float minHeight =  numeric_limits<float>::max();
          float maxHeight = -numeric_limits<float>::max();
          
          int i = ci*(chunkSize-1);
          int j = cj*(chunkSize-1);
          
          int ij = i+j*width;
          
          for(int vivj2=0,vj=0;vj<chunkSize;vj++,ij+=(width-chunkSize))
          {
            for(int vi=0;vi<chunkSize;vi++,vivj2+=2,ij++)
            {
              const float h = heightData[ij];
              
              const float h2 = (0.5f*(heightData[ij+offsetTable[vivj2+0]]+
                                      heightData[ij+offsetTable[vivj2+1]]));
              
              minHeight = min(minHeight,h);    
              maxHeight = max(maxHeight,h);    
              
              heightFieldBuffer[heightFieldBufferOffset+0] = h;
              heightFieldBuffer[heightFieldBufferOffset+1] = h2;            
              heightFieldBufferOffset += 2;
            }
          }
          
          chunkMinMaxHeight(ci,cj) = Vec2f(minHeight,maxHeight);  
        }     
      }
      
      glBindBuffer(GL_ARRAY_BUFFER,vbo);      
      glBufferSubData(GL_ARRAY_BUFFER,0,(chunksPerWidth*chunkSize)*(chunksPerDepth*chunkSize)*2*sizeof(float),heightFieldBuffer);                  
    }
    
    void draw(const Mat4x4f& modelview,const Mat4x4f& projection,float lodMult,float lodBias)
    {
      for(int j=0;j<chunksPerDepth;j++)
      {
        for(int i=0;i<chunksPerWidth;i++)
        {
          float x = float(i*(chunkSize-1)+(chunkSize-1)/2)/float(maxDim-1);
          float z = float(j*(chunkSize-1)+(chunkSize-1)/2)/float(maxDim-1);
          float y = 0.5f*(chunkMinMaxHeight(i,j)(0)+chunkMinMaxHeight(i,j)(1));
          
          Vec4f c = modelview*Vec4f(x,y,z,1.0f);
          chunkCenter(i,j) = Vec3f(c(0),c(1),c(2));
        }
      }

      GLint viewport[4];
      glGetIntegerv(GL_VIEWPORT,(GLint*)&viewport);
      
      int viewportWidth = viewport[2];
      int viewportHeight = viewport[3];
      
      float maxPixels = 8.0f; // !!
      
      float modelviewScale = norm(Mat3x3f(modelview(0,0),modelview(0,1),modelview(0,2),
                                          modelview(1,0),modelview(1,1),modelview(1,2),
                                          modelview(2,0),modelview(2,1),modelview(2,2))*Vec3f(1.0f/float(maxDim-1),0.0f,0.0f));
  
      Vec4f nearMinY = inverse(projection)*Vec4f(0.0f,-1.0f,-1.0f,1.0f);
      Vec4f nearMaxY = inverse(projection)*Vec4f(0.0f,+1.0f,-1.0f,1.0f);
      
      nearMinY = nearMinY / nearMinY(3);
      nearMaxY = nearMaxY / nearMaxY(3);
      
      float fovY = 2.0f*(atan((abs(nearMaxY(1)-nearMinY(1))/2.0f)/abs(nearMinY(2))));
      
      float projectionScale = float(viewportHeight) / tan(fovY/2.0f);
  
      for(int j=0;j<chunksPerDepth;j++)
      {
        for(int i=0;i<chunksPerWidth;i++)
        {        
          float lodDist = 1.0f;
          for(int lod=0;lod<numLods-1;lod++)
          {
            chunkMaxLodDist[lod](i,j) = lodBias+(lodDist*modelviewScale*projectionScale)/maxPixels;
            lodDist *= lodMult;
          }
        }
      }
      
      for(int j=0;j<chunksPerDepth;j++)
      {
        for(int i=0;i<chunksPerWidth;i++)
        {                            
          float chunkDist = norm(chunkCenter(i,j));
         
          int lod = 0;
  
          while(chunkDist > chunkMaxLodDist[lod](i,j) && lod<numLods-1)
          {
            lod++;
          }
  
          if (lod==numLods-1)
          {
            chunkLod(i,j) = float(lod);
          }
          else
          {
            float d0 = (lod>0) ? chunkMaxLodDist[lod-1](i,j) : 0.0f;
            float d1 = chunkMaxLodDist[lod](i,j);
            
            chunkLod(i,j) = float(lod) + (chunkDist-d0)/(d1-d0);
          }
        }
      }
      
      // TODO: chunkMaxLodDist fixup
      // TODO: frustum culling
      // TODO: front-to-back sorting
      
      GLint program;
      glGetIntegerv(GL_CURRENT_PROGRAM,&program);
      
      const GLint attribVertexY = glGetAttribLocation(program,"glhfVertexY");
      const GLint attribVertexMorphY = glGetAttribLocation(program,"glhfVertexMorphY");
      const GLint attribVertexBorderMask = glGetAttribLocation(program,"glhfVertexBorderMask");
      const GLint attribVertexLevel = glGetAttribLocation(program,"glhfVertexLevel");
      const GLint attribVertexXZ = glGetAttribLocation(program,"glhfVertexXZ");
  
      const GLint uniformChunkLod = glGetUniformLocation(program,"glhfChunkLod");
      const GLint uniformNeighborLods = glGetUniformLocation(program,"glhfNeighborLods");
      const GLint uniformChunkOriginXZ = glGetUniformLocation(program,"glhfChunkOriginXZ");
      const GLint uniformTexCoordScale = glGetUniformLocation(program,"glhfTexCoordScale");
      const GLint uniformTexCoordOffset = glGetUniformLocation(program,"glhfTexCoordOffset");

      Vec2f texCoordScale = Vec2f(float(maxDim-1)/float(width-1),float(maxDim-1)/float(depth-1))*Vec2f(1.0f-1.0f/float(width),1.0f-1.0f/float(depth));
      Vec2f texCoordOffset = 0.5f*Vec2f(1.0f/float(width),1.0f/float(depth));

      glUniform2f(uniformTexCoordScale,texCoordScale(0),texCoordScale(1));
      glUniform2f(uniformTexCoordOffset,texCoordOffset(0),texCoordOffset(1));

      glEnableVertexAttribArray(attribVertexY);
      glEnableVertexAttribArray(attribVertexMorphY);
      glEnableVertexAttribArray(attribVertexBorderMask);
      glEnableVertexAttribArray(attribVertexLevel);
      glEnableVertexAttribArray(attribVertexXZ);
  
      glBindBuffer(GL_ARRAY_BUFFER, vbo);
  
      glVertexAttribPointer(attribVertexBorderMask,4,GL_FLOAT,GL_FALSE,4*4,(void*)(borderMaskBufferOffset));
      glVertexAttribPointer(attribVertexLevel,1,GL_FLOAT,GL_FALSE,1*4,(void*)(vertexLevelBufferOffset));
      glVertexAttribPointer(attribVertexXZ,2,GL_FLOAT,GL_FALSE,2*4,(void*)(vertexCoordBufferOffset));
  
      for(int j=0;j<chunksPerDepth;j++)
      {
        for(int i=0;i<chunksPerWidth;i++)
        {
          int lod = (int)floorf(chunkLod(i,j));
          
          int stitchMask = geomipmap_stitch_mask(i>0                && ((int)floorf(chunkLod(i-1,j)))>lod,
                                              i<chunksPerWidth-1 && ((int)floorf(chunkLod(i+1,j)))>lod,
                                              j>0                && ((int)floorf(chunkLod(i,j-1)))>lod,
                                              j<chunksPerDepth-1 && ((int)floorf(chunkLod(i,j+1)))>lod);
          
          float originX = float(i*(chunkSize-1))/float(maxDim-1);
          float originZ = float(j*(chunkSize-1))/float(maxDim-1);
  
          float lodSelf = chunkLod(i,j);        
          
          float lodLeft   = (i>0)                ? max(lodSelf,chunkLod(i-1,j)) : lodSelf;
          float lodRight  = (i<chunksPerWidth-1) ? max(lodSelf,chunkLod(i+1,j)) : lodSelf;
          float lodTop    = (j>0)                ? max(lodSelf,chunkLod(i,j-1)) : lodSelf;
          float lodBottom = (j<chunksPerDepth-1) ? max(lodSelf,chunkLod(i,j+1)) : lodSelf;        
          
          glUniform1f(uniformChunkLod,lodSelf);
          glUniform4f(uniformNeighborLods,lodLeft,lodRight,lodTop,lodBottom);
          glUniform2f(uniformChunkOriginXZ,originX,originZ);
                  
          glBindBuffer(GL_ELEMENT_ARRAY_BUFFER,chunkIbo(lod,stitchMask));
  
          glVertexAttribPointer(attribVertexY,1,GL_FLOAT,GL_FALSE,2*sizeof(float),(void*)(((i+j*chunksPerWidth)*(chunkSize*chunkSize))*2*sizeof(float)));
          glVertexAttribPointer(attribVertexMorphY,1,GL_FLOAT,GL_FALSE,2*sizeof(float),(void*)(((i+j*chunksPerWidth)*(chunkSize*chunkSize))*2*sizeof(float)+1*sizeof(float)));
  
          glDrawElements(GL_TRIANGLES,3*geomipmap_num_tris(chunkSize,lod,stitchMask),GL_UNSIGNED_INT,0);      
        }
      }
    }
    
    void cleanup()
    {
      delete[] chunkMaxLodDist;
      
      delete[] heightFieldBuffer;
      
      delete[] offsetTable;
      
      glDeleteBuffers(1,&vbo);
  
      for(int lod=0;lod<numLods;lod++)
      {
        for(int stitchMask=0;stitchMask<16;stitchMask++)
        {
          glDeleteBuffers(1,&chunkIbo(lod,stitchMask));
        }
      }
    };
  
  private:
    int width;
    int depth;
    
    int chunkSize; 
    
    int chunksPerWidth;
    int chunksPerDepth;
  
    int numLods;
    
    int maxDim;
    
    Array2<GLuint> chunkIbo;
    GLuint vbo;
    
    int borderMaskBufferOffset;
    int vertexLevelBufferOffset;
    int vertexCoordBufferOffset;
    
    Array2<Vec3f>  chunkCenter;  
    Array2<float>* chunkMaxLodDist;  
    Array2<float>  chunkLod;
    Array2<Vec2f>  chunkMinMaxHeight;
    
    float* heightFieldBuffer;
    int*   offsetTable;
  };

  vector<GLHeightField*> glhfs;  
}

GLuint glhfCreateHeightfield(GLint width,GLint depth,GLint chunksize)
{
  GLHeightField* glhf = new GLHeightField();
  
  glhf->init(width,depth,chunksize);
  
  int heightfield = -1;
  
  for(int i=0;i<glhfs.size();i++)
  {
    if (glhfs[i]==0)
    {
      heightfield = i;
      break;
    }
  }
  
  if (heightfield==-1)
  {
    glhfs.push_back(0);
    heightfield = glhfs.size()-1;
  }
  
  glhfs[heightfield] = glhf;
    
  return (GLuint)heightfield;
}

void glhfHeightfieldData(GLuint heightfield,GLenum type,GLvoid* data)
{
  glhfs[heightfield]->update((float*)data);
}

void glhfDrawHeightfield(GLuint heightfield,GLfloat* modelview,GLfloat* projection)
{
  Mat4x4f matModelView = Mat4x4f(modelview[0],modelview[4],modelview[ 8],modelview[12],
                                 modelview[1],modelview[5],modelview[ 9],modelview[13],
                                 modelview[2],modelview[6],modelview[10],modelview[14],
                                 modelview[3],modelview[7],modelview[11],modelview[15]);

  Mat4x4f matProjection = Mat4x4f(projection[0],projection[4],projection[ 8],projection[12],
                                  projection[1],projection[5],projection[ 9],projection[13],
                                  projection[2],projection[6],projection[10],projection[14],
                                  projection[3],projection[7],projection[11],projection[15]);
                                 
  glhfs[heightfield]->draw(matModelView,matProjection,2.0,0.0);
}

void glhfDeleteHeightfield(GLuint heightfield)
{
  glhfs[heightfield]->cleanup();
  
  delete glhfs[heightfield];
  
  if (heightfield==glhfs.size()-1)
  {
    glhfs.pop_back();
  }
  else
  {
    glhfs[heightfield] = 0;
  }
}
